import torch
import torch.nn as nn
import torch.nn.functional as F

resnet_blocks = {
    'resnet50':[3, 4, 6, 3],
    'resnet101':[3, 4, 23, 3]
}

class bottle_neck(nn.Module):
    def __init__(self, input_dim, layer1_dim, stride=1):
        super(bottle_neck, self).__init__()

        self.expansion = 4
        self.main_block = nn.Sequential(
            nn.Conv2d(input_dim, layer1_dim, 1, stride=stride, padding=0),
            nn.BatchNorm2d(layer1_dim),
            nn.ReLU(),
            nn.Conv2d(layer1_dim, layer1_dim, 3, stride=1, padding=1),
            nn.BatchNorm2d(layer1_dim),
            nn.ReLU(),
            nn.Conv2d(layer1_dim, layer1_dim*self.expansion, 1, stride=1, padding=0),
            nn.BatchNorm2d(layer1_dim*self.expansion),
        )
    
        if stride != 1 or input_dim != layer1_dim*self.expansion:
            self.short_cut = nn.Sequential(
                nn.Conv2d(input_dim, layer1_dim*self.expansion, 1, stride=stride, padding=0),
                nn.BatchNorm2d(layer1_dim*self.expansion),
            )
        else:
            self.short_cut = nn.Sequential()

    def forward(self,input):
        out = F.relu(self.main_block(input) + self.short_cut(input))
        return out

class Resnet(nn.Module):
    def __init__(self, resner_deepth=50, dataset='imagenet'):
        super(Resnet,self).__init__()
        if resner_deepth > 34:
            self.layer_module = bottle_neck
            self.layer_exp = 4
            self.output_dim = 2048
        else:
            self.layer_exp = 1
            self.output_dim = 512

        self.input_dim = 64
        self.blocks = resnet_blocks['resnet'+str(resner_deepth)]

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(3, stride=2, padding=1),
        )

        convx = []
        for i in range(0,4):
            convx.append(self.create_layer(64*(2**i), self.layer_module, self.blocks[i], stride=(1 if i < 1 else 2)))
        self.conv2_5 = nn.Sequential(*convx)

        if dataset == 'imagenet':
            self.pool = nn.AvgPool2d(7,stride=1)
            self.fc = nn.Sequential(
                nn.Linear(self.output_dim,1000),
                nn.Softmax(),
            )
        else:
            self.pool = nn.AvgPool2d(1,stride=1)
            self.fc = nn.Sequential(
                nn.Linear(self.output_dim,10),
                nn.Softmax(),
            )

    def create_layer(self, layer1_dim=None, layer_module=None, n_blocks=None, stride=None):
        layer = []
        strides = [stride] + [1]*(n_blocks - 1)
        for stride in strides:
            layer.append(layer_module(self.input_dim, layer1_dim, stride))
            self.input_dim = layer1_dim * self.layer_exp

        return nn.Sequential(*layer)

    def forward(self, input):
        out = self.pool(self.conv2_5(self.conv1(input)))
        out = self.fc(out.view(out.size()[0],-1))
        return out


